import heapq
from collections import deque
from graphviz import Digraph
import heapq
from collections import deque

def visualize_search_tree(parent_map, initial, goal, filename='search_tree'):
    """
    可视化搜索树
    :param parent_map: 父节点映射字典 {子节点: 父节点}
    :param initial: 初始状态
    :param goal: 目标状态
    :param filename: 输出文件名
    """
    dot = Digraph(comment='Search Tree', format='png')
    dot.attr(rankdir='TB', nodesep='0.5', ranksep='0.5')

    # 添加所有节点
    all_nodes = set(parent_map.keys()) | set(parent_map.values()) | {initial}

    for node in all_nodes:
        # 将状态转换为3x3矩阵的字符串表示
        matrix = '\n'.join([' '.join(map(str, node[i:i + 3])) for i in range(0, 9, 3)])
        # 特殊样式标记初始状态和目标状态
        if node == initial:
            dot.node(str(node), label=matrix, shape='box', color='green', style='filled')
        elif node == goal:
            dot.node(str(node), label=matrix, shape='box', color='red', style='filled')
        else:
            dot.node(str(node), label=matrix, shape='box')

    # 添加边关系
    for child, parent in parent_map.items():
        dot.edge(str(parent), str(child))

    # 渲染并保存文件
    dot.render(filename, view=True)
def get_moves(state):
    blank = state.index(0)#获取0的索引
    row, col = blank // 3, blank % 3#获取行数和列数
    moves = []
    #上下左右移动空格，上即行-1
    directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]
    for dr, dc in directions:
        new_row, new_col = row + dr, col + dc
        if 0 <= new_row < 3 and 0 <= new_col < 3:#判断移动的位置是否合理
            new_blank = new_row * 3 + new_col#获取0位置的新索引
            new_state = list(state)#定义新状态变量，准备获取新的八数码棋状态（一个序列）
            new_state[blank], new_state[new_blank] = new_state[new_blank], new_state[blank]#进行新状态的赋值，即改变空格位置
            moves.append(tuple(new_state))#将新生成的八数码棋盘状态转换为元组（tuple），并添加到可能的移动列表（moves）中
    return moves

# 启发函数1：错位数（不计空格）
def h1(state, goal):
    return sum(1 for s, g in zip(state, goal) if s != g and s != 0)#zip的作用是把steate和goal的同一位置作为一对

# 启发函数2：曼哈顿距离（需预先计算目标位置字典）
def h2(state, goal_positions):
    distance = 0
    for idx in range(9):
        num = state[idx]
        if num == 0:
            continue
        target_row, target_col = goal_positions[num]
        current_row, current_col = idx // 3, idx % 3
        distance += abs(current_row - target_row) + abs(current_col - target_col)#计算要走到目标位置的步数
    return distance#每个棋子要走到目标位置的步数之和

def a_star(initial, goal, heuristic_func):
    open_heap = []#open_heap：优先队列，按f(n) = g(n) + h(n)排序，优先处理f小的节点。
    initial_h = heuristic_func(initial)# 用估价函数计算初始状态的启发值（如错位数或曼哈顿距离）
    heapq.heappush(open_heap, (initial_h + 0, 0, initial))#将元组(initial_h + 0, 0, initial)=(f+g,g,state)加入堆
    #将一个元素推入堆（open_heap），并根据元素的第一个值（f值）=估计代价+实际代价来自动调整堆结构，保证堆顶元素始终是最小的f值。
    best_g = {initial: 0}#best_g：字典，记录每个节点最小的g(n)实际代价（避免重复处理高代价路径）。
    parent_map = {}#parent_map：字典，记录每个节点的父节点（用于回溯路径）
    generated = 1#实际处理（计算f=h+g）的节点数初始值
    expanded = 0

    while open_heap:
        current_f, current_g, current_state = heapq.heappop(open_heap)#弹出f最小的状态
        if current_state == goal:
            path = []#如果已达成目标状态，返回移动步数和节点数
            while current_state in parent_map:
                path.append(current_state)
                current_state = parent_map[current_state]
            path.append(initial)
            path.reverse()
            return path, expanded, generated, parent_map  # 增加返回parent_map

        if current_g > best_g.get(current_state, float('inf')):
            continue##如果当前状态的`g`值不是最优的（即存在更优的路径已经记录在`best_g`中），跳过该节点。否则，扩展当前状态，生成所有可能的邻居状态。

        expanded += 1#扩展节点+1，准备扩展当前状态的邻居状态

        for neighbor in get_moves(current_state):#生成所有可能的邻居状态。
            new_g = current_g + 1#对于每个邻居状态，计算邻居状态的`g`值（`current_g + 1`）。
            if neighbor not in best_g or new_g < best_g.get(neighbor, float('inf')):#如果新`g`值更优，则更新`best_g`，计算新的`f`值，将邻居状态推入堆，并记录父状态
                best_g[neighbor] = new_g
                new_h = heuristic_func(neighbor)#计算估价函数的值
                new_f = new_g + new_h
                heapq.heappush(open_heap, (new_f, new_g, neighbor))
                parent_map[neighbor] = current_state
                generated += 1#生成节点＋1

    return None, expanded, generated,parent_map  # 增加返回parent_map

def bfs(initial, goal):
    queue = deque([(initial, 0)])#记录初始状态步数为0
    visited = set()
    visited.add(initial)#访问，即加入初始状态节点进已处理节点
    parent_map = {}
    generated = 1#实际处理的节点
    expanded = 0#初始，扩展的领居状态为0个

    while queue:
        current_state, current_g = queue.popleft()#取出队列中第一个节点，取出它的具体状态，实际代价
        expanded += 1#准备扩展当前状态的子节点，所以当前状态是一个扩展节点
        if current_state == goal:
            path = []
            while current_state in parent_map:
                path.append(current_state)
                current_state = parent_map[current_state]
            path.append(initial)
            path.reverse()
            return path, expanded, generated

        for neighbor in get_moves(current_state):#获取当前状态的邻居状态进行处理
            if neighbor not in visited:#如果邻居状态未处理过
                visited.add(neighbor)
                parent_map[neighbor] = current_state
                queue.append((neighbor, current_g + 1))#将邻居状态和该状态的实际代价加入队列
                generated += 1#邻居状态处理完毕，实际处理节点数+1

    return None, expanded, generated
def reversenum(node):
    """计算状态对应的逆序数,奇偶性一致则有解"""
    Sum = 0
    for i in range(1,9):
        num = 0
        for j in range(0,i):
          if node[j]>node[i] and node[i] != 0:
              num = num + 1
        Sum += num
    return Sum

def main():
    # 初始状态和目标状态
    initial_state = (2, 8, 3, 1, 6, 4, 7, 0, 5)
    goal_state = (1, 2, 3, 8, 0, 4, 7, 6, 5)
    # 判断从初始状态是否可以达到目标状态
    if (reversenum(initial_state) % 2) != (reversenum(goal_state) % 2):
        print("该目标状态不可达！")#奇偶性不一致
        return
    else:
        print("可解：")
    # 定义h1的启发函数（直接比较目标状态）
    h1_func = lambda state: h1(state, goal_state)

    # 定义h2的启发函数（预先计算目标位置）
    goal_positions = {num: (i // 3, i % 3) for i, num in enumerate(goal_state) if num != 0}
    # 将原来数字的序列位置转化为（行号，列号）位置，//为整数除法。
    h2_func = lambda state: h2(state, goal_positions)


    # 修改后
    path_h1, expanded_h1, generated_h1, parent_h1 = a_star(initial_state, goal_state, h1_func)
    path_h2, expanded_h2, generated_h2, parent_h2 = a_star(initial_state, goal_state, h2_func)

    # path_bfs, expanded_bfs, generated_bfs = bfs(initial_state, goal_state)
    visualize_search_tree(parent_h1, initial_state, goal_state, 'astar_tree1')
    visualize_search_tree(parent_h2, initial_state, goal_state, 'astar_tree2')
    # 输出结果
    print("A* with h1 (错位数):")
    print(f"扩展节点数: {expanded_h1}, 生成节点数: {generated_h1}, 移动步数: {len(path_h1) - 1}")
    # for i in range(len(path_h1)):
    #     state = path_h1[i]
    #     print(f"Step {i}:")
    #     for row in range(0, 9, 3):  # 每3个元素为一组
    #         print(" ".join(map(str, state[row:row + 3])))  # 将数字转换为字符串并用空格连接
    #     print()  # 添加空行分隔每一步
    print("\nA* with h2 (曼哈顿距离):")
    print(f"扩展节点数: {expanded_h2}, 生成节点数: {generated_h2}, 移动步数: {len(path_h2) - 1}")
    #
    # print("\nBFS:")
    # print(f"扩展节点数: {expanded_bfs}, 生成节点数: {generated_bfs}, 移动步数: {len(path_bfs) - 1}")

if __name__ == "__main__":
    main()